#!/usr/bin/env python

import sys
import os
import gtk
import gobject
import gtkmozembed
import getopt

def usage():
    print "Usage: %s [-u url]\r\n" % sys.argv[0]
    print "Options: --url -u URL             The url\r\n"
    print "         --output -o OUTPUT       Output file"
    print "         --height HEIGHT          Height in pixels"
    print "         --width WIDTH            Width in pixels"
    print "         --timeout SECONDS        Timout"
    print "         --delay SECONDS          Delay"
    print "         --verbose -v             Very noisy"
    print "         --help -h                This usage"
    print "         --xvfb -x                Run in Xvfb"
    sys.exit(1)
"""
Take png images of web sites. Renders web data with gtkmozembed. Uses timer and
counter, reading data from gtk.MozEmbed signal prototypes into a list. When the counter
has encountered the interval of empty of data, the page has finished loading,
and the image is taken. 

This script is specific to the Linux environment, and can make use of Xvfb for DISPLAY.  
Happy web-page-imaging!

Posted: Sun May 23 19:50:01 PDT 2010 v1.0 --karl@webmedianow.com
"""
 
class Moz2Png:
    def __init__(self, URL = None, WIDTH = None, HEIGTH = None, OUTPUT = None):
        # Kick off the timer
        self.timer = gobject.timeout_add(1000, self.loadtimer, self)
        # Initialize gtkmozembed
        self.moz = gtkmozembed.MozEmbed()
        win = gtk.Window(gtk.WINDOW_TOPLEVEL)
        win.add(self.moz)
        win.show_all()
        # Load the url
        self.moz.load_url(URL)
        self.moz.set_chrome_mask(gtkmozembed.FLAG_WINDOWRESIZEON) # Removes the scroll bar
        self.moz.set_size_request(WIDTH, HEIGTH)
 
        # Connect ALL the signals - http://www.pygtk.org/pygtkmozembed/class-gtkmozembed.html
        self.moz.connect('net-start', self.moz_net_start)
        self.moz.connect('progress', self.moz_progress)
        self.moz.connect('net-state', self.moz_net_state)
        self.moz.connect('net-stop', self.moz_net_stop)
        self.moz.connect('js-status', self.moz_js_status)
        self.moz.connect('visibility', self.moz_visibility)
        self.moz.connect('open-uri', self.moz_open_uri)
        self.moz.connect('title', self.moz_title)
        self.moz.connect('location', self.moz_location)
        self.moz.connect('link-message', self.moz_link_message)
        self.moz.connect('new-window', self.moz_new_window)
        self.moz.connect('destroy-browser', self.moz_destroy_browser)

    def moz_title(self, data = None):
        list.insert(0, "title")

    def moz_location(self, data = None):
        list.insert(0, "location")

    def moz_link_message(self, data = None):
        list.insert(0, "link-message")

    def moz_new_window(self, retval, chromemask, data = None):
        list.insert(0, "new-window")

    def moz_destroy_browser(self, data = None):
        list.insert(0, "destroy-browser")
        gtk.main_quit() # This shouldn't occur...

    def moz_visibility(self, visibility, date = None):
        vis_num = len(list)
        visibility = "visibilit%d" % vis_num
        list.insert(0, visibility)
        # http://www.pygtk.org/pygtkmozembed/class-gtkmozembed.html
        # visibility :	 if TRUE then the window should be shown. If it's FALSE it should be hidden.
        # The "visibility" signal is emitted when the toplevel window in question needs to be shown or hidden. 
        return True

    def moz_open_uri(self, uri, date = None):
        uri_num = len(list)
        open_uri = "open-uri%d" % uri_num
        list.insert(0, open_uri)
        # http://www.pygtk.org/pygtkmozembed/class-gtkmozembed.html
        # The "open-uri" signal is emitted when when the document tries to load a new document
        # If you return FALSE the new document will be loaded.
        # By returning TRUE you are saying "don't load this document."
        return False

    def moz_net_start(self, data = None):
        list.insert(0, "net-start")

    def moz_progress(self, data, cur, max):
        list.insert(0, cur)

    def moz_net_state(self, data, cur, max):
        list.insert(0, cur)

    def moz_js_status(self, data = None):
        js_num = len(list)
        js_status = "js-status%d" % js_num
        list.insert(0, js_status)

    def moz_net_stop(self, data = None):
        list.insert(0, "net-stop")

    def loadtimer(self, object):

        if len(round) == 0:
          print "Timeout %d is up!" % timeout
          self.snapshot()

        if self.delayer() >= delay:
          if verbose == True: print "Delay: " , list
          self.snapshot()

        if self.loadtimer > 0:
          list.insert(0, 0)
  
        if verbose == True: print "Round %s" % len(round) , list
        round.pop() 
        return True

    def delayer(self):
        count = 0
        for index, item in enumerate(list[:delay]):
          if item == list[0]:
            count += 1
        return count

    def snapshot(self, data = None):
        if verbose == True: print "Taking SnapShot... "
        sys.stdout.flush()
        (x, y, width, height, depth) = self.moz.window.get_geometry()
        pixbuf = gtk.gdk.Pixbuf(gtk.gdk.COLORSPACE_RGB, False, 8, width, height)
        pixbuf.get_from_drawable(self.moz.window, self.moz.get_colormap(), 0, 0, 0, 0, width, height)
        pixbuf.save(OUTPUT, 'png')    # Write the png file
        print "Done: %s" % OUTPUT
        gtk.main_quit()

def run_xvfb():
    server_num = int(os.getpid() + 1e6) # Start 'xvfb' instance by replacing the current process
    newArgs = ["xvfb-run", "--auto-servernum", "--server-num", str(server_num), "--server-args=-screen 0, %dx%dx24" % (WIDTH, HEIGHT), sys.argv[0]]
    skipArgs = 0
    for i in range(1, len(sys.argv)):
        if skipArgs > 0:
            skipArgs -= 1
        elif sys.argv[i] in ["-x", "--xvfb"]:
            skipArgs = 0
        else:
            newArgs.append(sys.argv[i])
    if verbose == True: print "Executing %s" % " ".join(newArgs)
    os.execvp(newArgs[0],newArgs[1:])

 
if __name__ == "__main__":

    URL = ''
    WIDTH = 1024
    HEIGHT = 768
    OUTPUT = 'out.png'
    delay = 3
    timeout = 15
    verbose = False
    xvfb = False

    try:
        opts, args = getopt.getopt(sys.argv[1:], "ho:vu:x", ["height=", "width=", "help", "output=",
                                                             "url=", "delay=", "timeout=", "xvfb"])
    except getopt.GetoptError, err:
        print str(err) # will print something like "option -a not recognized"
        usage() # print help information and exit:
    for o, a in opts: 
        if o in ("-h", "--help"):
            usage()
            sys.exit(2)
        elif o in ("-v", "--verbose"):
            verbose = True
        elif o in ("-o", "--output"): 
            OUTPUT = a
        elif o in ("-u", "--url"):
            URL = a
        elif o in ("--width"):
            WIDTH = int(a)
        elif o in ("--height"): 
            HEIGHT = int(a)
        elif o in ("--delay"):
            delay = int(a)
        elif o in ("--timeout"):
            timeout = int(a)
        elif o in ("-x", "--xvfb"):
            xvfb = True
        else:
            assert False, "unhandled option"

    if URL == '': 
        usage() # print help information and exit:

    if xvfb == True:
       run_xvfb()
       
    list = [1]
    round = [0]*timeout
 
    Moz2Png(URL, WIDTH, HEIGHT, OUTPUT)
    gtk.main()


